﻿#include "Plate2D.h"
#include "Application.h"
#include "Mesh.h"
#include "Timer.h"
#include "Transient.h"
#include <vtkm/Bounds.h>
#include <vtkm/Geometry.h>

struct MeshModifier_Plate2D : vtkm::worklet::WorkletCellNeighborhood
{
  using ControlSignature = void(CellSetIn,
                                FieldInNeighborhood centroid,
                                FieldOut blk_id,
                                FieldOut bnd_id);
  using ExecutionSignature = void(Boundary, _2, _3, _4);

  template<typename NeighIn, typename BndIdType>
  VTKM_EXEC void operator()(const vtkm::exec::BoundaryState& boundary,
                            const NeighIn& centroid,
                            vtkm::Int16& blk_id,
                            BndIdType& bnd_id) const
  {
    blk_id = 0;
    bnd_id = { -1, -1, -1, -1 };
    if (boundary.IJK[0] == 0) // left
    {
      bnd_id[3] = 3;
    }
    if (boundary.IJK[0] == boundary.PointDimensions[0] - 1) // right
    {
      bnd_id[1] = 1;
    }
    if (boundary.IJK[1] == boundary.PointDimensions[1] - 1) // top
    {
      bnd_id[2] = 2;
    }
    if (boundary.IJK[1] == 0) // bottom
    {
      Vec3f p = centroid.Get(0, 0, 0);
      if (p[0] < 1. / 6)
        bnd_id[0] = 0;
      else
        bnd_id[0] = 4;
    }
  }

  VTKM_EXEC bool IsCellValid(Vec3f p) const
  {
    return Square(p);
    //return Sphere(p);
    //return Naca0012(p);
    //return ForwardStep(p);
  }

  VTKM_EXEC bool Square(const Vec3f& p) const
  {
    Real x = p[0];
    Real y = p[1];
    Real a = 0.1;
    if (x < a && x > -a && y < a && y > -a)
    {
      return false;
    }
    else if (x < 0.4 && x > 0.3 && y < 0.7 && y > 0.6)
    {
      return false;
    }
    else if (x < 0.8 && x > 0.6 && y < 0.4 && y > 0.2)
    {
      return false;
    }
    else if (x < 1.5 && x > 1.3 && y < 0.2 && y > -0.1)
    {
      return false;
    }
    else
    {
      return true;
    }
  }

  VTKM_EXEC bool Sphere(const Vec3f& p) const
  {
    if (vtkm::Magnitude(p) < 0.1)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 0.5, 0.2, 0 }) < 0.2)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 0.5, -0.2, 0 }) < 0.2)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 0.7, 0.2, 0 }) < 0.15)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 0.7, -0.2, 0 }) < 0.15)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 0.9, 0.3, 0 }) < 0.15)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 0.9, -0.3, 0 }) < 0.15)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 1.7, -0.3, 0 }) < 0.15)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 1.5, 0.5, 0 }) < 0.15)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 2.5, 0.3, 0 }) < 0.15)
    {
      return false;
    }
    if (vtkm::Magnitude(p - Vec3f{ 2.5, -0.3, 0 }) < 0.3)
    {
      return false;
    }

    return true;
  }

  VTKM_EXEC bool Naca0012(const Vec3f& p) const
  {
    Real x = p[0];
    Real y = p[1];

    Real a0 = +0.2969, a1 = -0.1260, a2 = -0.3516, a3 = +0.2843, a4 = -0.1036;
    Real t = 0.12;

    if (x < 0 || x > 1)
    {
      return true;
    }
    else
    {
      Real yt =
        5 * t * (a0 * vtkm::Sqrt(x) + a1 * x + a2 * x * x + a3 * x * x * x + a4 * x * x * x * x);
      if (y > yt || y < -yt)
      {
        return true;
      }
      else
      {
        return false;
      }
    }
  }

  VTKM_EXEC bool ForwardStep(const Vec3f& p) const
  {
    vtkm::Bounds box(Vec3f{ 0.6, 0, 0 }, Vec3f{ 3, 0.2, 0 });
    if (box.Contains(p))
    {
      return false;
    }
    return true;
  }
};



struct Plate2DIC : WorkletInitCondition<Plate2DIC>
{
  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    FloatDefault gam = 1.4, gamm1 = gam - 1, mach = 0.2;
    FloatDefault u = 1;
    FloatDefault v = 0;
    FloatDefault density = 1;
    FloatDefault pre = 1. / (gam * mach * mach);

    ValueType value;
    value[0] = density;
    value[1] = density * u;
    value[2] = density * v;
    value[3] = density * 0;
    value[4] = pre / gamm1 + 0.5 * density * (u * u + v * v);

    return value;
  }
};




struct ComputeResidual_Plate2D : vtkm::worklet::WorkletCellNeighborhood
{
  VTKM_CONT
  ComputeResidual_Plate2D(FloatDefault t)
    : _t(t)
  {
  }

  using ControlSignature = void(CellSetIn,
                                FieldInNeighborhood solution,
                                FieldOut Rhs,
                                FieldIn centroid,
                                FieldIn spacing,
                                FieldIn blk_id,
                                FieldIn bnd_id);
  using ExecutionSignature = void(Boundary, _2, _3, _4, _5, _6, _7);

  template<typename NeighIn, typename RhsType>
  VTKM_EXEC void operator()(const vtkm::exec::BoundaryState& boundary,
                            const NeighIn& value,
                            RhsType& rhs,
                            const Vec3f& centroid,
                            const Vec3f& spacing,
                            const vtkm::Int16& blk_id,
                            const Vec<vtkm::Int16, 4>& bnd_id) const
  {
    if (blk_id == -1)
    {
      rhs[0] = 0;
      rhs[1] = 0;
      rhs[2] = 0;
      rhs[3] = 0;
      rhs[4] = 0;
      return;
    }

    FloatDefault penalty = 4;
    Vec3f normal{ 0, 0, 0 };
    DataPack<RhsType::NUM_COMPONENTS> dl, dr, lift;
    dl.uh = value.Get(0, 0, 0);
    dl.InvisTerm();
    {
      normal = Vec3f{ -1, 0, 0 };
      dr.uh = value.Get(-1, 0, 0);
      RhsType lflux, rflux;
      Vec3f p = centroid;
      p[0] = p[0] - 0.5 * spacing[0];
      BoundaryCondition(bnd_id[3], _t, p, normal, dl, dr);

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      dr.InvisTerm();
      RiemannFlux(dl, dr, normal, lflux);
      lift.ViscousTerm();

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        lflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }


      dr.uh = value.Get(1, 0, 0);
      normal = Vec3f{ 1, 0, 0 };
      p = centroid;
      p[0] = p[0] + 0.5 * spacing[0];
      BoundaryCondition(bnd_id[1], _t, p, normal, dl, dr);

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      dr.InvisTerm();
      RiemannFlux(dl, dr, normal, rflux);
      lift.ViscousTerm();
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        rflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      rhs = (rflux + lflux) / spacing[0];
    }

    {
      dr.uh = value.Get(0, -1, 0);
      RhsType lflux, rflux;
      normal = Vec3f{ 0, -1, 0 };
      Vec3f p = centroid;
      p[1] = p[1] - 0.5 * spacing[1];
      BoundaryCondition(bnd_id[0], _t, p, normal, dl, dr);
      dr.InvisTerm();
      RiemannFlux(dl, dr, normal, lflux);

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }
      lift.ViscousTerm();

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        lflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      dr.uh = value.Get(0, 1, 0);
      normal = Vec3f{ 0, 1, 0 };
      p = centroid;
      p[1] = p[1] + 0.5 * spacing[1];
      BoundaryCondition(bnd_id[2], _t, p, normal, dl, dr);
      dr.InvisTerm();
      RiemannFlux(dl, dr, normal, rflux);

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }
      lift.ViscousTerm();

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        rflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      rhs += (rflux + lflux) / spacing[1];
    }
  }

  template<typename DataPackType>
  VTKM_EXEC void BoundaryCondition(const vtkm::Int16& bid,
                                   const FloatDefault& t,
                                   const Vec3f& p,
                                   const Vec3f& normal,
                                   const DataPackType& dl,
                                   DataPackType& dr) const
  {
    if (bid == 1) // 出口
    {
      auto& ul = dl.uh;
      auto& ur = dr.uh;
      ur[0] = ul[0];
      ur[1] = ul[1];
      ur[2] = ul[2];
      ur[3] = ul[3];
      ur[4] = ul[4];
    }
    if (bid == 3 || bid == 0 || bid == 2) // 进口
    {
      FloatDefault gam = 1.4, gamm1 = gam - 1, mach = 0.2;
      FloatDefault u = 1;
      FloatDefault v = 0;
      FloatDefault density = 1;
      FloatDefault pre = 1. / (gam * mach * mach);

      auto& ur = dr.uh;
      ur[0] = density;
      ur[1] = density * u;
      ur[2] = density * v;
      ur[3] = density * 0;
      ur[4] = pre / gamm1 + 0.5 * density * (u * u + v * v);
    }
    if (bid == 4) //wall
    {
      auto& ul = dl.uh;
      auto& ur = dr.uh;
      Vec3f mom{ ul[1], ul[2], ul[3] };
      Real vn = vtkm::Dot(mom, normal);
      ur[0] = ul[0];
      ur[1] = 0;
      ur[2] = 0;
      ur[3] = 0;
      ur[4] = ul[4];
    }
  }

  template<typename DataPackType, typename VType>
  VTKM_EXEC void RiemannFlux(const DataPackType& dl,
                             const DataPackType& dr,
                             const Vec3f& normal,
                             VType& riemann_flux) const
  {
    for (IdComponent eq = 0; eq < VType::NUM_COMPONENTS; eq++)
    {
      riemann_flux[eq] =
        vtkm::Dot(dl.inv_term[eq] + dr.inv_term[eq], normal) / 2 + 1.5 * (dl.uh[eq] - dr.uh[eq]);
    }
  }


  FloatDefault _t;
};

RegisterObject(Plate2D);
Plate2D::Plate2D(const Configuration& cfg)
  : NavierStokesEquation(cfg)
{
}

void Plate2D::Init()
{
  NavierStokesEquation::Init();
  vtkm::cont::Invoker invoker;
  const auto& centroid = _mesh.GetCellField("center");

  vtkm::cont::Invoker{}(MeshModifier_Plate2D{}, _mesh.GetCellSet(), centroid, _blk_id, _bnd_id);
}

void Plate2D::InitialCondition()
{
  vtkm::cont::Timer timer;
  timer.Start();
  vtkm::cont::Invoker invoker;
  invoker(Plate2DIC{}, _mesh.GetCellSet(), _solution, _mesh.GetCoordinate());

  console::Success("Plate2D初始条件设置完成", timer.GetElapsedTime(), "s");
}

void Plate2D::UpdateResidual()
{
  auto& transient = static_cast<Transient&>(*(_app.GetExecutioner()));
  Real current_time = transient.CurrentTime();

  Timer timer;
  timer.Start();
  const auto& centroid = _mesh.GetCellField("center");
  const auto& spacing = _mesh.GetCellField("spacing");
  vtkm::cont::Invoker{}(ComputeResidual_Plate2D{ current_time },
                        _mesh.GetCellSet(),
                        _solution,
                        _rhs,
                        centroid,
                        spacing,
                        _blk_id,
                        _bnd_id);
  console::Info("Plate2D::UpdateResidual运行完成，", "时间: ", timer.GetElapsedTime(), " s");
}